import os
import re
import numpy as np
import matplotlib.pyplot as plt
import argparse
from pathlib import Path

def process_files(files):
    cum = []

    # Process each file in sorted order
    for file_path in files:
        if os.path.isfile(file_path):
            print("file_path_valid:", file_path)
            with open(file_path, 'r') as file:
                # Column indices
                column_index_1 = 1
                column_index_2 = 5

                # Arrays for storing data
                mar_array = []
                times_array = []
                data_started = False

                for line in file:
                    if not data_started:
                        if line.startswith("record#"):
                            data_started = True
                        continue
                    
                    if data_started and line.strip():
                        row = re.split(r'\s+', line.strip())
                        if len(row) > column_index_1:
                            mar_array.append(row[column_index_1])
                        if len(row) > column_index_2:
                            times_array.append(row[column_index_2])

                # Convert to numpy arrays
                MAR = np.array(mar_array)
                TIMES = np.array(times_array).astype(int)

                # Get indices for "MAR" and "OFL"
                indices_mar = np.where(MAR == 'MAR')[0]
                indices_ofl = np.where(MAR == 'OFL')[0]

                # Calculate how many "OFL" entries come before each "MAR"
                ofl_counts_before_mar = np.searchsorted(indices_ofl, indices_mar)

                # Adjust MAR indices based on OFL counts
                indices_mar_selected = indices_mar - ofl_counts_before_mar - np.arange(indices_mar.shape[0])

                # Corresponding times for the selected MAR indices
                TIMES_mar = TIMES[indices_mar_selected]

                # Split the TIMES array into slices based on MAR indices
                times_slice = []
                slices = np.split(TIMES, indices_mar_selected)
                for i, Slices in enumerate(slices):
                    if i >= 1:
                        times_slice.append(Slices)

                # Adjust the last slice to match the average photon count
                temp = sum([times_slice[i].shape[0] for i in range(indices_mar.shape[0] - 1)])
                mean_photons = temp // (indices_mar.shape[0] - 1)
                times_slice[-1] = times_slice[-1][:mean_photons]

                # Denoise by filtering photon counts within the range [245, 265]
                times_slice_denoisied = []
                sum_counts = np.zeros(indices_mar.shape[0])
                for i in range(indices_mar.shape[0]):
                    temp = times_slice[i][(times_slice[i] > 245) & (times_slice[i] < 265)]
                    sum_counts[i] = len(temp)
                    times_slice_denoisied.append(temp)
                
                # Collect cumulative photon counts
                cum.extend(sum_counts)
                print("processed the file:", file_path)

    return cum

def plot_data(cum, output_path, output_name):
    # Plot the cumulative photon counts
    plt.plot(cum)
    plt.title("Number of photons corresponding to patterns")
    plt.xlabel("Index")
    plt.ylabel("Photons")

    # Ensure the directory exists
    output_dir = Path(output_path)
    output_dir.mkdir(parents=True, exist_ok=True)

    # Save the figure to the specified path with the specified name
    output_file = output_dir / f"{output_name}.png"
    plt.savefig(output_file)
    print(f"Image saved to: {output_file}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Process a list of files and generate plots.")
    parser.add_argument('files', nargs='+', help='List of file paths to process')
    parser.add_argument('--output-path', type=str, required=True, help='Path where the image will be saved')
    parser.add_argument('--output-name', type=str, required=True, help='Name of the output image file (without extension)')

    args = parser.parse_args()
    file_list = args.files
    output_path = args.output_path
    output_name = args.output_name

    print(f"Received file list from command line: {file_list}")
    print(f"Output image will be saved at: {output_path}/{output_name}.png")

    cum_data = process_files(file_list)
    plot_data(cum_data, output_path, output_name)